








































import numpy as np
from sortedcollections import SortedDict
from quant.utils import logging


class MarketData:
    def __init__(self, event, exchange, symbol, frequency, routing_key, recv_time, server_time, data_id, server_id, data):
        self.event = event
        self.exchange = exchange
        self.symbol = symbol
        self.frequency = frequency

        self.routing_key = routing_key
        self.recv_time = recv_time
        self.server_time = server_time
        self.data_id = data_id
        self.server_id = server_id
        self.data = data  # 1.OrderBook(); 2.list() of trades:[['buy', 32000, 0.1], [...]]; 3.ticker((bp, bv), (ap, av))

    def __repr__(self):
        result = '{} {} {}'.format(self.event, self.exchange, self.symbol)
        result += ' recv:{} server:{}'.format(self.recv_time, self.server_time)
        result += ' id:{} server:{}'.format(self.data_id, self.server_id)
        return result

    def describe(self):
        result = '{} {} {} {}'.format(self.event, self.exchange, self.symbol, self.frequency)
        return result


class OrderBook:
    def __init__(self):
        self._buy_sorted_dic = SortedDict()
        self._sell_sorted_dic = SortedDict()
        self._side_to_dic = {'buy': self._buy_sorted_dic, 'sell': self._sell_sorted_dic}

    def __repr__(self):
        result = '{} | {}'.format([[p, v] for p, v in self.items('buy')], [[p, v] for p, v in self.items('sell')])
        return result

    def book_repr(self, count=5):
        bids = []
        asks = []
        for p, v in self.items('buy'):
            bids.append((p, v))
            if len(bids) > count:
                break
        for p, v in self.items('sell'):
            asks.append((p, v))
            if len(asks) > count:
                break
        # bids = [[p, v] for p, v in self.items('buy')][:count]
        # asks = [[p, v] for p, v in self.items('sell')][:count]
        result = '\n'.join([f'{p} {v}' for p, v in reversed(asks)])
        result += '\n------------\n'
        result += '\n'.join([f'{p} {v}' for p, v in bids])
        return result

    def __getitem__(self, key):
        return self._side_to_dic[key[0]][key[1]]

    def __setitem__(self, key, value):
        # self._build_id()

        if isinstance(key, str):
            dic = self._side_to_dic[key]
            dic.clear()
            dic.update(value)

        elif isinstance(key, tuple):
            side, price = key
            if value > 0:
                self._side_to_dic[side][price] = value
            else:
                try:
                    del self._side_to_dic[side][price]
                except KeyError:
                    pass

        self._cut_len()

    def __len__(self):
        return len(self._buy_sorted_dic) + len(self._sell_sorted_dic)

    def clear(self):
        self._buy_sorted_dic.clear()
        self._sell_sorted_dic.clear()
        # self._build_id()

    def check_cross(self, error_name):
        # self._build_id()
        try:
            buy_1, _ = self._buy_sorted_dic.peekitem(-1)
            sell_1, _ = self._sell_sorted_dic.peekitem(0)
        except IndexError:
            return

        if buy_1 > sell_1:
            del self._buy_sorted_dic[buy_1]
            del self._sell_sorted_dic[sell_1]
            logging.error(f'{error_name} buy_1({buy_1}) cross with sell_1({sell_1})')

    def item(self, side, depth):
        if depth == 0:
            return None, None
        if side == 'buy' or side == 'bid':
            i = -depth
            dic = self._buy_sorted_dic
        else:
            i = depth-1
            dic = self._sell_sorted_dic

        try:
            return dic.peekitem(i)
        except IndexError:
            return None, None

    def items(self, side):
        if side == 'sell' or side == 'ask':
            yield from self._sell_sorted_dic.items()
        else:
            for price in reversed(self._buy_sorted_dic):
                yield price, self._buy_sorted_dic[price]

    def get_middle(self):
        try:
            buy_1, _ = self._buy_sorted_dic.peekitem(-1)
            sell_1, _ = self._sell_sorted_dic.peekitem(0)
        except IndexError:
            return None
        return (buy_1 + sell_1) / 2

    def has_price(self, side, price):
        return price in self._side_to_dic[side]

    def list(self, side, count):
        if side == 'buy':
            li = list(self._buy_sorted_dic.items())
            li = li[-count:]
            li.reverse()
        else:
            li = list(self._sell_sorted_dic.items())
            li = li[:count]

        return li

    def _cut_len(self):
        while len(self._buy_sorted_dic) > BOOK_MAX_LEN:
            self._buy_sorted_dic.popitem(0)
        while len(self._sell_sorted_dic) > BOOK_MAX_LEN:
            self._sell_sorted_dic.popitem(-1)


class OrderBookNew(OrderBook):  # trying
    def copy(self):
        new = OrderBookNew()
        new._buy_sorted_dic = self._buy_sorted_dic.copy()
        new._sell_sorted_dic = self._sell_sorted_dic.copy()
        new._side_to_dic = {'buy': new._buy_sorted_dic, 'sell': new._sell_sorted_dic}
        return new

    def diff(self, pre):
        pre_bid = pre._buy_sorted_dic
        this_bid = self._buy_sorted_dic
        keys = set(pre_bid)
        keys.update(this_bid)
        bid_diff = [(k, this_bid.get(k, 0) - pre_bid.get(k, 0)) for k in keys]

        pre_ask = pre._sell_sorted_dic
        this_ask = self._sell_sorted_dic
        keys = set(pre_ask)
        keys.update(this_ask)
        ask_diff = [(k, this_ask.get(k, 0) - pre_ask.get(k, 0)) for k in keys]

        bid_diff = [v for v in bid_diff if v[1] != 0]
        ask_diff = [v for v in ask_diff if v[1] != 0]
        return bid_diff, ask_diff


BOOK_MAX_LEN = 400


if __name__ == '__main__':
    from quant.markets import OrderBook
    from quant.exchanges.basics import check_book_cross
    from quant.utils import perf_test, set_test_mode
    from utils import build_test_book
    set_test_mode()

    def perf_book_update():
        book = OrderBook()

        no_order_update = [[i, i] for i in range(600)]

        def perf():
            book['buy'] = no_order_update

        perf_test(perf)


    book = build_test_book(count=400)

    def fn():
        li = book.list('buy', 200)
        arr[:] = li

    def fn1():
        book.list('sell', 200)


    book = OrderBook()
    bids = [
        [1, 1],
        [3, 3],
        [2, 2],
    ]
    book['buy'] = bids
    book['sell'] = bids

    print(book)























